"""
算法描述:
knn分类器首先对一组已标记(已知)的人脸进行训练，然后可以预测这个人
在实时流中找到k张最相似的脸(在欧几里得距离下，人脸特征接近的图像)
在它的训练集，并执行多数投票(可能加权)对他们的标签。


$ pip3 install scikit-learn
$ pip3 install numpy
$ pip3 install opencv-contrib-python

"""

import cv2
import math
from sklearn import neighbors
import os
import os.path
import pickle
from PIL import Image, ImageDraw
import face_recognition
from face_recognition.face_recognition_cli import image_files_in_folder
import numpy as np
# import base64

# 图片编码为Base64
# with open('图片路径','rb') as f:
#     a=base64.b64encode(f.read())
# print(a)

# Base64解码为图片
# b=base64.b64decode('base64编码')
# with open('新的图片路径','wb') as f:
#     f.write(b)

ALLOWED_EXTENSIONS = [".jpg", ".png", ".jpeg", ".webp"]


def train(train_dir, model_save_path=None, n_neighbors=None, knn_algo='ball_tree', verbose=False):
    """
    Trains a k-nearest neighbors classifier for face recognition.

    :param train_dir: directory that contains a sub-directory for each known person, with its name.

     (View in source code to see train_dir example tree structure)

     Structure:
        <train_dir>/
        ├── <person1>/
        │   ├── <somename1>.jpeg
        │   ├── <somename2>.jpeg
        │   ├── ...
        ├── <person2>/
        │   ├── <somename1>.jpeg
        │   └── <somename2>.jpeg
        └── ...

    :param model_save_path: 存储模型的路径
    :param n_neighbors: (optional) number of neighbors to weigh in classification. Chosen automatically if not specified
    :param knn_algo: 训练knn模型，所使用的算法（默认：ball_tree）
    :param verbose: 图片中的人脸不止一张时，是否训练该图片（默认只训练图片中只有一张人脸的图片）
    :return: 返回训练的结果
    """
    print('开始模型训练')

    X = []
    y = []

    # 遍历每个人的目录
    for class_dir in os.listdir(train_dir):
        if not os.path.isdir(os.path.join(train_dir, class_dir)):
            continue

        # 训练每个人的目录下的图片
        for img_path in image_files_in_folder(os.path.join(train_dir, class_dir)):
            # 过滤非图片格式的文件
            if os.path.splitext(img_path[img_path.rfind('/'): len(img_path)])[1].lower() not in ALLOWED_EXTENSIONS:
                continue

            print(f"处理图片:{img_path}")

            image = face_recognition.load_image_file(img_path)
            # number_of_times_to_upsample:在寻找人脸时将原始图像放大n倍，可以帮助在图像中找到可能被遗漏的较小的面孔，但是会导致运行得更慢。
            face_bounding_boxes = face_recognition.face_locations(image, number_of_times_to_upsample=1, model="hog")

            if len(face_bounding_boxes) != 1:
                # If there are no people (or too many people) in a training image, skip the image.
                if verbose:
                    print("Image {} not suitable for training: {}".format(img_path, "Didn't find a face" if len(
                        face_bounding_boxes) < 1 else "Found more than one face"))
            else:
                # Add face encoding for current image to the training set
                X.append(face_recognition.face_encodings(image, known_face_locations=face_bounding_boxes)[0])
                y.append(class_dir)

    # Determine how many neighbors to use for weighting in the KNN classifier
    if n_neighbors is None:
        n_neighbors = int(round(math.sqrt(len(X))))
        if verbose:
            print("Chose n_neighbors automatically:", n_neighbors)

    # Create and train the KNN classifier
    # n_neighbors：用于kneighbors查询的默认邻居的数量
    # algorithm：计算最近邻居用的算法（默认为 ‘auto’）：
    # # ‘ball_tree’ 是为了克服kd树高纬失效而发明的，其构造过程是以质心C和半径r分割样本空间，每个节点是一个超球体。
    # # ‘kd_tree’ 构造kd树存储数据以便对其进行快速检索的树形数据结构，kd树也就是数据结构中的二叉树。以中值切分构造的树，每个结点是一个超矩形，在维数小于20时效率高。
    # # ‘brute’ 使用暴力搜索.也就是线性扫描，当训练集很大时，计算非常耗时
    # # ‘auto’ 会基于传入fit方法的内容，选择最合适的算法。
    # weights：权重(默认为 ‘uniform’)：
    # # ‘uniform’ : 统一的权重. 在每一个邻居区域里的点的权重都是一样的。
    # # ‘distance’ : 权重点等于他们距离的倒数。使用此函数，更近的邻居对于所预测的点的影响更大。
    # # [callable] : 一个用户自定义的方法，此方法接收一个距离的数组，然后返回一个相同形状并且包含权重的数组。

    knn_clf = neighbors.KNeighborsClassifier(n_neighbors=n_neighbors, algorithm=knn_algo, weights='distance')
    knn_clf.fit(X, y)

    print('模型训练结束')

    # Save the trained KNN classifier
    if model_save_path is not None:
        with open(model_save_path, 'wb') as f:
            pickle.dump(knn_clf, f)

    print('模型储存成功')

    return knn_clf


def predict(X_frame_BGR, knn_clf=None, model_path=None, distance_threshold=0.5, trait_model='hog', image_size_multiple=1):
    """
    Recognizes faces in given image using a trained KNN classifier

    :param trait_model: 识别人脸用的特征(hog,cnn)
    :param X_frame_BGR: 处理的人脸的多维数组
    :param knn_clf: (optional) a knn分类器对象，如果它不存，则knn_clf文件路径必须存在
    :param model_path: 识别模型的保存路径，必须为knn_clf文件.
    :param distance_threshold: 欧几里德距离的阈值：阈值太低容易造成无法成功识别人脸，太高容易造成人脸识别混淆（原始默认值：0.5，这里针对亚洲人改为0.4）
    :return: 一个list对象，包含人名，人脸边框，人脸轮廓坐标: [(name, bounding box, face landmark) , ...].
            如果人脸未知，则以'unknown'命名
    """

    if knn_clf is None and model_path is None:
        raise Exception("Must supply knn classifier either thourgh knn_clf or model_path")

    # Load a trained KNN model (if one was passed in)
    if knn_clf is None:
        with open(model_path, 'rb') as f:
            knn_clf = pickle.load(f)

    # OpenCV使用的是BGR模式，而PIL使用的是RGB模式。
    # Only 'RGB' (8-bit RGB, 3 channels) and 'L' (black and white)

    X_face_locations = face_recognition.face_locations(X_frame_BGR, number_of_times_to_upsample=image_size_multiple, model=trait_model)
    X_face_landmarks = face_recognition.face_landmarks(X_frame_BGR, X_face_locations)  # 面部轮廓位置

    # If no faces are found in the image, return an empty result.
    if len(X_face_locations) == 0:
        print('图片中没有人脸')
        return []

    # Find encodings for faces in the test image
    faces_encodings = face_recognition.face_encodings(X_frame_BGR, known_face_locations=X_face_locations)

    # Use the KNN model to find the best matches for the test face
    closest_distances = knn_clf.kneighbors(faces_encodings, n_neighbors=1)
    are_matches = [closest_distances[0][i][0] <= distance_threshold for i in range(len(X_face_locations))]

    # Predict classes and remove classifications that aren't within the threshold
    # [('obama', (204, 932, 590, 546), {'chin': [(525, 348), (533, 405), (542, 461), (560, 513), (590, 560), (631, 598), (678, 626), (732, 644), (787, 643), (830, 626), (851, 590), (869, 549), (882, 507), (888, 462), (890, 417), (889, 373), (883, 330)], 'left_eyebrow': [(598, 323), (625, 296), (662, 278), (705, 276), (745, 284)], 'right_eyebrow': [(783, 282), (812, 269), (844, 266), (873, 275), (885, 301)], 'nose_bridge': [(769, 315), (775, 345), (782, 374), (790, 405)], 'nose_tip': [(738, 439), (761, 441), (784, 443), (801, 437), (817, 431)], 'left_eye': [(644, 336), (666, 323), (689, 319), (710, 330), (691, 337), (667, 340)], 'right_eye': [(800, 323), (820, 309), (840, 309), (857, 320), (842, 326), (821, 326)], 'top_lip': [(710, 519), (739, 499), (767, 487), (783, 491), (798, 486), (818, 492), (835, 506), (826, 504), (799, 497), (784, 501), (767, 500), (722, 515)], 'bottom_lip': [(835, 506), (820, 522), (804, 530), (788, 533), (770, 534), (743, 530), (710, 519), (722, 515), (769, 511), (786, 512), (801, 507), (826, 504)]})]

    return [(pred, loc, landmark) if rec else ("unknown", loc, landmark) for pred, loc, rec, landmark in zip(knn_clf.predict(faces_encodings), X_face_locations, are_matches, X_face_landmarks)]


def show_prediction_labels_on_image(frame, predictions, deal_pattern='image', save_face_detection='false', save_img_path='', face_landmark_open='false'):
    """
    Shows the face recognition results visually.

    :param deal_pattern:处理的模式（image/video）,如果时video会对图片进行缩放，以加速识别
    :param frame: 需要绘制的图片
    :param predictions: predict函数的返回结果，包括人名、人脸位置、人脸特征点
    :param save_face_detection:保存识别出来的裁剪后的人脸
    :param save_img_path:裁剪后的图片存储路径
    :param face_landmark_open:是否开启勾画面部特征轮廓
    :return opencv suited image to be fitting with cv2.imshow fucntion: 返回绘制的图像
    """

    pil_image = Image.fromarray(frame)
    draw = ImageDraw.Draw(pil_image)

    save_path_index = 1  # 用保存裁剪用的图片编号
    for name, (top, right, bottom, left), face_landmarks in predictions:
        # enlarge the predictions for the full sized image.
        if deal_pattern == 'video':
            top *= 4
            right *= 4
            bottom *= 4
            left *= 4
        else:
            pass

        # Draw a box around the face using the Pillow module
        draw.rectangle(((left, top), (right, bottom)), outline=(0, 0, 255))

        if save_face_detection.lower() == 'true' and len(save_img_path) != 0:  # 保存才剪出来的图片
            image = Image.open(save_img_path)

            save_path = save_img_path[0:save_img_path.rfind('.')] + str(save_path_index) + '.jpg'
            print(f"人脸裁剪结果的保存路径:{save_path}")

            save_img = image.crop((left, top, right, bottom))

            save_img.save(save_path)
            save_img.close()
            save_path_index += 1
        else:
            pass

        # There's a bug in Pillow where it blows up with non-UTF-8 text
        # when using the default bitmap font
        name = name.encode("UTF-8")

        # Draw a label with a name below the face
        text_width, text_height = draw.textsize(name)
        draw.rectangle(((left, bottom - text_height - 10), (right, bottom)), fill=(0, 0, 255), outline=(0, 0, 255))
        draw.text((left + 6, bottom - text_height - 5), name, fill=(255, 255, 255, 255))

        # {'chin': [(525, 348), (533, 405), (542, 461), (560, 513), (590, 560), (631, 598), (678, 626), (732, 644), (787, 643), (830, 626), (851, 590), (869, 549), (882, 507), (888, 462), (890, 417), (889, 373), (883, 330)],
        # 'left_eyebrow': [(598, 323), (625, 296), (662, 278), (705, 276), (745, 284)],
        # 'right_eyebrow': [(783, 282), (812, 269), (844, 266), (873, 275), (885, 301)],
        # 'nose_bridge': [(769, 315), (775, 345), (782, 374), (790, 405)],
        # 'nose_tip': [(738, 439), (761, 441), (784, 443), (801, 437), (817, 431)],
        # 'left_eye': [(644, 336), (666, 323), (689, 319), (710, 330), (691, 337), (667, 340)],
        # 'right_eye': [(800, 323), (820, 309), (840, 309), (857, 320), (842, 326), (821, 326)],
        # 'top_lip': [(710, 519), (739, 499), (767, 487), (783, 491), (798, 486), (818, 492), (835, 506), (826, 504), (799, 497), (784, 501), (767, 500), (722, 515)],
        # 'bottom_lip': [(835, 506), (820, 522), (804, 530), (788, 533), (770, 534), (743, 530), (710, 519), (722, 515), (769, 511), (786, 512), (801, 507), (826, 504)]}

        if face_landmark_open.lower() == 'false':
            continue
        else:
            pass

        for key, marks in face_landmarks.items():  # 绘制面部轮廓点
            for i in range(len(marks)):
                point = marks[i]
                p0 = point[0]
                p1 = point[1]
                if deal_pattern == 'video':
                    p0 *= 4
                    p1 *= 4
                else:
                    pass

                draw.ellipse(((p0, p1), (p0 + 1, p1 + 1)), fill=(0, 255, 0), outline=(0, 255, 0), width=0)

    # Remove the drawing library from memory as per the Pillow docs.
    del draw
    # Save image in open-cv format to be able to show it.

    opencvimage = np.array(pil_image)
    return opencvimage


def cameraCatchComparison(known_img_model_file_path):
    """
    使用摄像头进行人脸识别，速度比较慢（慢3-5s）
    :param known_img_model_file_path: 已知模型的人脸识别
    :return: 无
    """
    # 计算机摄像设备索引
    CAMERA_IDX = 0
    # 调用摄像头，conf.CAMERA_IDX为摄像头索引，默认为0，也可以这样写cap = cv2.VideoCapture(0)
    cap = cv2.VideoCapture(CAMERA_IDX)
    while True:
        # 读取一帧数据
        if cap.isOpened() == False:
            break
        ok, frame = cap.read()

        # Convert the image from BGR color (which OpenCV uses) to RGB color (which face_recognition uses)
        # OpenCV是BGR格式，PIL是RGB格式
        # rgb_frame = frame
        rgb_frame = frame[:, :, ::-1]

        if not ok:
            break

        img = cv2.resize(rgb_frame, (0, 0), fx=0.25, fy=0.25)

        predictions = predict(img, model_path=known_img_model_file_path)

        frame = show_prediction_labels_on_image(frame, predictions, deal_pattern='video')

        # 显示图像
        cv2.imshow('Face Recognition Video', frame)

        # 监听输入，按q退出
        if cv2.waitKey(1) & 0xFF == ord('q'):
            break

    # 释放摄像头并销毁所有窗口
    cap.release()
    cv2.destroyAllWindows()


def faceDetectionComparison(known_img_model_file_path, comparison_img_path, distance_threshold=0.5, image_size_multiple=1, trait_model='hog'):
    """
    图片对照
    :param known_img_model_file_path:模型存储路径
    :param comparison_img_path:对比的人物图片
    :param distance_threshold:距离阈值
    :param image_size_multiple:图片尺寸
    :param trait_model:提取的面部特征
    :return:返回图片中包含的已知人名的list
    """

    img = cv2.imread(comparison_img_path)
    # ('biden', (226, 464, 315, 374)， )

    predictions = predict(X_frame_BGR=img, model_path=known_img_model_file_path, image_size_multiple=image_size_multiple, distance_threshold=distance_threshold, trait_model=trait_model)

    # 不包含unknow
    # predict_people = list(set(filter(lambda x: x != 'unknown', [prediction[0] for prediction in predictions])))

    # 包含unkonw，并计数
    predict_people = list(set([prediction[0] if prediction[0] != 'unknown' else prediction[0] + str(index) for index, prediction in enumerate(predictions)]))

    return predict_people



def drawImageFace(known_img_model_file_path, img_path, save_face_detection='false', face_landmark_open='false', distance_threshold=0.45, image_size_multiple=1, trait_model='hog'):
    """
    绘制图片上的人脸
    :param known_img_model_file_path:模型的存储路径
    :param img_path:图片路径
    :param save_face_detection：是否保存出裁剪出来的人脸
    :param face_landmark_open:是否开启人脸轮廓的识别
    :param image_size_multiple:图片的大小倍数，用于寻找较小的面孔
    :return:无
    """
    img = cv2.imread(img_path)

    predictions = predict(img, model_path=known_img_model_file_path, image_size_multiple=image_size_multiple, distance_threshold=distance_threshold, trait_model=trait_model)

    frame = show_prediction_labels_on_image(img, predictions, deal_pattern='image', save_face_detection=save_face_detection, save_img_path=img_path, face_landmark_open=face_landmark_open)

    cv2.imshow('Face Recognition Image', frame)

    # OpenCV使用的是BGR模式，而PIL使用的是RGB模式。
    # frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB) # 如果不转格式，直接保存后，图片是蓝色的（bgr）
    # im = Image.fromarray(frame)
    # im.save("res.jpg")

    cv2.waitKey(0)


def main():
    import sys
    argv_len = len(sys.argv)
    if argv_len < 2:
        print(f"Please run: {sys.argv[0]} train|recognition /example/img.jpg")
        exit()

    model_path = os.path.dirname(os.path.abspath(sys.argv[0])) + '/trained_knn_model.clf'
    print("model_path:{}".format(model_path))
    # 根据命令行参数选择操作
    operation = sys.argv[1]
    if operation == "train":
        classifier = train("knn_examples/train", model_save_path=model_path, n_neighbors=2)
    elif operation == "recognition":
        if argv_len < 3:
            print('请输入有效的图片路径')
        else:
            # # 简约版
            comparison_img_path = sys.argv[2]
            print(faceDetectionComparison(model_path, comparison_img_path))
    else:
        raise ValueError(f"Unsupported operation: {operation}")


if __name__ == "__main__":

    main()

    # #  训练模型
    # classifier = train("knn_examples/train", model_save_path="trained_knn_model.clf", n_neighbors=2)

    # #  测试图片
    # known_img_model_file_path = '/Users/mac/PycharmProjects/faceRecognitionProject/faceRecognition/trained_knn_model.clf'
    # img_path = 'knn_examples/test/5.jpg'
    # # 简约版
    # print(faceDetectionComparison(known_img_model_file_path, img_path))
    # # 详细版
    # print(faceDetectionComparison(known_img_model_file_path, img_path, distance_threshold=0.4, image_size_multiple=1, trait_model='hog'))

    # #  绘制图片上的人脸
    # known_img_model_file_path = '/Users/mac/PycharmProjects/faceRecognitionProject/faceRecognition/trained_knn_model.clf'
    # img_path = '/Users/mac/Downloads/人脸识别/xjp/xjp2.jpeg'
    # # # 简约版
    # drawImageFace(known_img_model_file_path, img_path)
    # # # 详细版
    # drawImageFace(known_img_model_file_path, img_path,  save_face_detection='false', face_landmark_open='true', distance_threshold=0.4, image_size_multiple=2, trait_model='hog')

    # #  测试视频
    # known_img_model_file_path = '/Users/mac/PycharmProjects/faceRecognitionProject/faceRecognition/trained_knn_model.clf'
    # cameraCatchComparison(known_img_model_file_path)

    pass
